﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

using System.Diagnostics.Contracts;

using SimpleIR;
using PInfoModel;

namespace PInfoModel.Analyses
{
	public class IDMapAtCallInfoForFullProgram
	{
		public readonly List<IDMapAtCallInfoForAllInMethod> AllMethodIDMappings;

		private readonly CompleteAnalysisInfoForProgram m_pip;

		private readonly string m_dumpdir;

		private IDMapAtCallInfoForFullProgram(string dumpdir, CompleteAnalysisInfoForProgram pip)
		{
			this.m_pip = pip;
			this.m_dumpdir = dumpdir;

			this.AllMethodIDMappings = new List<IDMapAtCallInfoForAllInMethod>();
		}

		public static IDMapAtCallInfoForFullProgram MakeIDMappingForProgram(string dumpdir, CompleteAnalysisInfoForProgram pip)
		{
			IDMapAtCallInfoForFullProgram idmcfp = new IDMapAtCallInfoForFullProgram(dumpdir, pip);

			foreach(KeyValuePair<IRInvokeIdentity, CompleteAnalysisInfoMethod> cmc in pip.MethodInfo)
			{
				IDMapAtCallInfoForAllInMethod idmc = new IDMapAtCallInfoForAllInMethod(cmc.Key);
				foreach(CompleteAnalysisInfoMethodContext cmic in cmc.Value.AllCallContexts)
					idmc.ProcessSingleContext(pip, cmic, pip.Program.MethodBodyTable.GetBodyForInvokeExactSigMatch(cmc.Key), dumpdir);

				idmcfp.AllMethodIDMappings.Add(idmc);
			}

			return idmcfp;
		}

		/// <summary>
		/// Write the info for this particular call mapping to the xml stream.
		/// </summary>
		public void XMLExport(System.Xml.XmlWriter xw)
		{
			xw.WriteStartElement("IDMapAtCallInfoFullProgram");
			xw.WriteAttributeString("MethodsCount", this.AllMethodIDMappings.Count.ToString());

			foreach(IDMapAtCallInfoForAllInMethod mmap in this.AllMethodIDMappings)
				mmap.XMLExport(xw);

			this.m_pip.Program.CallGraph.XMLExportSimple(xw, this.m_pip.Program);

			xw.WriteEndElement();
		}

		public void XMLDump()
		{
			System.Xml.XmlWriterSettings xws = new System.Xml.XmlWriterSettings() { CloseOutput = true, Indent = true };
			string filename = System.IO.Path.Combine(this.m_dumpdir, "callmappings.output");

			using(System.IO.FileStream fs = new System.IO.FileStream(filename, System.IO.FileMode.Create))
			{
				using(System.Xml.XmlWriter xw = System.Xml.XmlWriter.Create(fs, xws))
				{
					this.XMLExport(xw);
				}
			}
		}
	}

	public class IDMapAtCallInfoForAllInMethod
	{
		/// <summary>
		/// The invoke identity that this information is for
		/// </summary>
		public readonly IRInvokeIdentity ForMethod;

		/// <summary>
		/// Map from call-context ids to a list of bytecodes in the method with calls and the relevant call-info.
		/// </summary>
		public readonly Dictionary<int, List<IDMapAtCallInfo>> CallSiteToInfoMappings;

		/// <summary>
		/// Map from call-context ids to a list of memory locations that are visible at method entry.
		/// </summary>
		public readonly Dictionary<int, List<Tuple<IRType, MemLocID>>> MemLocsAtEntry;

		public IDMapAtCallInfoForAllInMethod(IRInvokeIdentity fmethod)
		{
			this.ForMethod = fmethod;
			this.CallSiteToInfoMappings = new Dictionary<int, List<IDMapAtCallInfo>>();
			this.MemLocsAtEntry = new Dictionary<int, List<Tuple<IRType, MemLocID>>>();
		}

		internal void ProcessSingleContext(CompleteAnalysisInfoForProgram pip, CompleteAnalysisInfoMethodContext cmmc, IRBodyDef bdf, string outputdir)
		{
			//compute at entry memory info
			List<Tuple<IRType, MemLocID>> mlids = new List<Tuple<IRType, MemLocID>>();
			foreach(PIMNode nn in cmmc.InState.Nodes)
			{
				List<MemLocID> mlidl = new List<MemLocID>();
				PIModel.AddAllLocationIDSForNodeToSet(nn, mlidl);

				foreach(MemLocID mlid in mlidl)
					mlids.Add(new Tuple<IRType, MemLocID>(nn.LocationType, mlid));
			}

			this.MemLocsAtEntry.Add(cmmc.ContextID, mlids);

			//compute the call site mapping info
			List<IDMapAtCallInfo> idmcl = new List<IDMapAtCallInfo>();
			foreach(IRBasicBlock bb in bdf.Body.BasicBlockList)
			{
				foreach(IRByteCode tbc in bb.Operations)
				{
					if(tbc is IRCallOp || tbc is IRAllocationOp)
					{
						IDMapAtCallInfo idmc = new IDMapAtCallInfo(outputdir, pip, bdf, cmmc, tbc);
						idmc.ResolveInfo();

						idmcl.Add(idmc);
					}
				}
			}

			this.CallSiteToInfoMappings.Add(cmmc.ContextID, idmcl);
		}

		/// <summary>
		/// Write the info for this particular call mapping to the xml stream.
		/// </summary>
		public void XMLExport(System.Xml.XmlWriter xw)
		{
			xw.WriteStartElement("IDMapAtCallInfoForMethod");
			xw.WriteAttributeString("MethodID", this.ForMethod.ToString());
			xw.WriteAttributeString("ContextCount", this.MemLocsAtEntry.Count.ToString());

			foreach(int ctxid in this.MemLocsAtEntry.Keys)
			{
				xw.WriteStartElement("SingleCallContext");
				xw.WriteAttributeString("ContextID", ctxid.ToString());

				xw.WriteStartElement("MemIDsAtEntry");
				xw.WriteAttributeString("Count", this.MemLocsAtEntry[ctxid].Count.ToString());
				foreach(Tuple<IRType, MemLocID> mlidp in this.MemLocsAtEntry[ctxid])
				{
					xw.WriteStartElement("EntryLocation");
					xw.WriteAttributeString("LocationType", mlidp.Item1.ToString());
					mlidp.Item2.XMLExport(xw);
					xw.WriteEndElement();
				}
				xw.WriteEndElement();

				xw.WriteStartElement("CallMapInfos");
				xw.WriteAttributeString("Count", this.CallSiteToInfoMappings[ctxid].Count.ToString());
				foreach(IDMapAtCallInfo idci in this.CallSiteToInfoMappings[ctxid])
					idci.XMLExport(xw);
				xw.WriteEndElement();

				xw.WriteEndElement();
			}

			xw.WriteEndElement();
		}
	}

	/// <summary>
	/// Class to implement code that takes a call with a line number and extracts the info on which memory region identifiers 
	/// are mapped into/out of each callee + dump heap graphs for inspection.
	/// </summary>
	public class IDMapAtCallInfo
	{
		/// <summary>
		/// Where to write the dgml files.
		/// </summary>
		private string m_dumpdir;

		/// <summary>
		/// The whole program info structure.
		/// </summary>
		private CompleteAnalysisInfoForProgram m_pip;

		/// <summary>
		/// The caller method body.
		/// </summary>
		private IRBodyDef m_body;

		/// <summary>
		/// The call bytecode we are interested in.
		/// </summary>
		private IRByteCode m_forbc;

		/// <summary>
		/// The caller context that we are doing the extraction for. In practice we would need to do it for all of them and summarize.
		/// </summary>
		private CompleteAnalysisInfoMethodContext m_cctx;

		////////////////////////////

		/// <summary>
		/// The heap state before the method is called.
		/// </summary>
		private PIModel m_prebytecodeexecstate;

		/// <summary>
		/// The heap state after the method returns.
		/// </summary>
		private PIModel m_postbytecodeexecstate;

		/// <summary>
		/// The list of all resolved call targets.
		/// </summary>
		private List<IDMapAtCallInfoSingleCallTarget> m_targets;

		////////////////////////////

		public IDMapAtCallInfo(string outdir, CompleteAnalysisInfoForProgram pip, IRBodyDef bdf, CompleteAnalysisInfoMethodContext cmctx, IRByteCode tbc)
		{
			this.m_dumpdir = outdir;
			this.m_pip = pip;
			this.m_body = bdf;
			this.m_cctx = cmctx;
			this.m_forbc = tbc;

			this.m_targets = new List<IDMapAtCallInfoSingleCallTarget>();
		}

		public static IDMapAtCallInfo DoInfoExtract(CompleteAnalysisInfoForProgram pip, IRBodyDef bdf, int line, string dumpdir)
		{
			//get the context 
			CompleteAnalysisInfoMethod cmm = pip.MethodInfo[bdf.InvokeID];

			if(cmm.AllCallContexts.Count > 1)
				Console.WriteLine("Method has multiple call contexts -- Just doing things for the first one!!");

			CompleteAnalysisInfoMethodContext cmic = cmm.AllCallContexts[0];

			//get the instruction
			IRBasicBlock tbb = bdf.Body.BasicBlockList.FirstOrDefault((IRBasicBlock bb) => bb.Operations.Any((IRByteCode tbc) => tbc.UseModLoc == line));
			if(tbb == null)
			{
				Console.WriteLine("Instruction not found!!");
				return null;
			}

			IRByteCode bc = tbb.Operations.FirstOrDefault((IRByteCode tbc) => tbc.UseModLoc == line);
			if(bc == null || (!(bc is IRCallOp) && !(bc is IRAllocationOp)))
			{
				Console.WriteLine("Instruction not found!!");
				return null;
			}

			IDMapAtCallInfo idimm = new IDMapAtCallInfo(dumpdir, pip, bdf, cmic, bc);
			idimm.ResolveInfo();

			return idimm;
		}

		internal void ResolveInfo()
		{
			////
			//get the caller state info
			Dictionary<int, ProgramPointInfo> ppimap = this.m_cctx.ProgramPointToPointInfoMap();
			if(!ppimap.ContainsKey(this.m_forbc.UseModLoc) || ppimap[this.m_forbc.UseModLoc].PreOperationModel == null)
				return;

			////
			//get the info on the pre/post state within the caller
			ProgramPointInfo ppi = ppimap[this.m_forbc.UseModLoc];
			this.m_prebytecodeexecstate = ppi.PreOperationModel;
			this.m_postbytecodeexecstate = ppi.PostOperationModel;

			if(this.m_forbc is IRAllocationOp)
			{
				IRType alloctype = ((IRAllocationOp)this.m_forbc).AllocType;
				if(alloctype is IRContainerType && ((IRContainerType)alloctype).IsArrayContainer)
					return;
			}

			List<IRInvokeIdentity> tiivl;
			if(this.m_forbc is IRCallOp)
				tiivl = new List<IRInvokeIdentity>(((IRCallOp)this.m_forbc).AnalysisResolvedCallTargets);
			else
				tiivl = new List<IRInvokeIdentity>() { ((IRAllocationOp)this.m_forbc).ConsIdentity };

			foreach(IRInvokeIdentity iiv in tiivl)
			{
				IRBodyDef calleeb = this.m_pip.Program.MethodBodyTable.GetBodyForInvokeExactSigMatch(iiv);
				
				if(calleeb.IsBuiltin)
				{
					IDMapAtCallInfoSingleCallTarget cttm = new IDMapAtCallInfoSingleCallTarget(calleeb, -1, null, null);
					this.m_targets.Add(cttm);
				}
				else
				{
					////
					//Lookup the info that relates the method call on the given line to the to callee methods and contexts that it may resolve to
					HashSet<CallerCalleeIDMapping> fccmapping = this.m_cctx.CallToCalleeStateMaps[this.m_forbc.UseModLoc];
					CallerCalleeIDMapping ccmapping = fccmapping.First((CallerCalleeIDMapping ccim) => ccim.CallerContextID == this.m_cctx.ContextID);

					////
					//get the info on the callee entry/exit heap structure
					CompleteAnalysisInfoMethodContext calleem = this.m_pip.LookupInfoForMethodAndContext(ccmapping.CalleeInvokeIdentity, ccmapping.CalleeContextID);
					IDMapAtCallInfoSingleCallTarget cttm = new IDMapAtCallInfoSingleCallTarget(calleeb, calleem.ContextID, calleem.InState, calleem.OutState);


					//just pick first mapping for example
					foreach(KeyValuePair<MemLocID, HashSet<MemLocID>> mmkvp in ccmapping.CallerIDToCalleeIDMap)
					{
						if(mmkvp.Key.IsHomeLocationId)
							continue;

						cttm.FromCallerToCallee.Add(mmkvp.Key, new HashSet<MemLocID>(mmkvp.Value));
					}

					if(calleem.OutState != null)
					{
						foreach(KeyValuePair<PIEscapeSite, HashSet<PIEscapeSite>> esch in calleem.CalleeEscToCallerEscMap)
						{
							MemLocID kid = MemLocID.MakeInternalSourceLocation(esch.Key.EscIdentity, esch.Key.EscLine, null);
							HashSet<MemLocID> escs = new HashSet<MemLocID>();

							foreach(PIEscapeSite tesci in esch.Value)
								escs.Add(MemLocID.MakeInternalSourceLocation(tesci.EscIdentity, this.m_forbc.UseModLoc, null));

							cttm.FromCalleeToCaller.Add(kid, escs);
						}
					}

					this.m_targets.Add(cttm);
				}
			}
		}

		public void PresentInfo()
		{
			if(this.m_prebytecodeexecstate == null)
			{
				Console.WriteLine("Call Never Executed.");
				return;
			}

			if(this.m_forbc is IRAllocationOp)
				Console.WriteLine("\nAllocation with ID=" + MemLocID.MakeInternalSourceLocation(this.m_forbc.UseModLoc, this.m_forbc.UseModLoc, null));

			//write the pre/post models
			string dgmlpre = this.m_prebytecodeexecstate.DGMLDumpPModel(this.m_forbc.UseModLoc);
			System.IO.File.WriteAllText(System.IO.Path.Combine(this.m_dumpdir, "beforestmt.dgml"), dgmlpre);

			string dgmlpost = this.m_postbytecodeexecstate.DGMLDumpPModel(this.m_forbc.UseModLoc);
			System.IO.File.WriteAllText(System.IO.Path.Combine(this.m_dumpdir, "afterstmt.dgml"), dgmlpost);

			if(this.m_forbc is IRAllocationOp)
			{
				IRType alloctype = ((IRAllocationOp)this.m_forbc).AllocType;
				if(alloctype is IRContainerType && ((IRContainerType)alloctype).IsArrayContainer)
				{
					Console.WriteLine("Is Array Allocation.");
					return;
				}
			}

			foreach(IDMapAtCallInfoSingleCallTarget idmst in this.m_targets)
			{
				if(idmst.CalleeTarget.IsBuiltin)
				{
					Console.WriteLine("Call To Builtin Method: " + idmst.CalleeTarget.InvokeID.ToString());
					return;
				}

				string dgmlentry = idmst.CalleeEntryState.DGMLDumpPModel(-1);
				System.IO.File.WriteAllText(System.IO.Path.Combine(this.m_dumpdir, "calleeenter.dgml"), dgmlentry);

				string dgmlexit = idmst.CalleeExitState.DGMLDumpPModel(-1);
				System.IO.File.WriteAllText(System.IO.Path.Combine(this.m_dumpdir, "calleeexit.dgml"), dgmlexit);

				//print out the mapping info

				Console.WriteLine("\nMapping from caller identifiers to callee identifiers:");
				foreach(KeyValuePair<MemLocID, HashSet<MemLocID>> mmset in idmst.FromCallerToCallee)
				{
					string mll = BuildIDMapString(mmset.Key, mmset.Value);
					Console.WriteLine(mll);
				}

				Console.WriteLine("\nMapping from callee identifiers (interal allocs) to caller identifiers:");
				foreach(KeyValuePair<MemLocID, HashSet<MemLocID>> mmset in idmst.FromCalleeToCaller)
				{
					string mll = BuildIDMapString(mmset.Key, mmset.Value);
					Console.WriteLine(mll);
				}
			}
		}

		/// <summary>
		/// Write the info for this particular call mapping to the xml stream.
		/// </summary>
		public void XMLExport(System.Xml.XmlWriter xw)
		{
			xw.WriteStartElement("CallIDMapForBC");
			xw.WriteAttributeString("CCIOffset", this.m_forbc.SourceLocationInfo.CCIOffset.ToString());
			xw.WriteAttributeString("NotExecuted", (this.m_prebytecodeexecstate == null).ToString());
			xw.WriteAttributeString("NeverReturns", (this.m_postbytecodeexecstate == null).ToString());

			bool isarrayalloc = false;
			if(this.m_forbc is IRAllocationOp)
			{
				IRType alloctype = ((IRAllocationOp)this.m_forbc).AllocType;
				if(alloctype is IRContainerType && ((IRContainerType)alloctype).IsArrayContainer)
					isarrayalloc = true;
			}

			xw.WriteAttributeString("Constructor", (this.m_forbc is IRAllocationOp).ToString());
			xw.WriteAttributeString("ArrayAllocation", isarrayalloc.ToString());

			if(this.m_forbc is IRAllocationOp)
			{
				xw.WriteAttributeString("AllocationType", ((IRAllocationOp)this.m_forbc).AllocType.ToString());
				xw.WriteStartElement("AllocatedLocation");
				MemLocID.MakeInternalSourceLocation(this.m_forbc.UseModLoc, this.m_forbc.UseModLoc, null).XMLExport(xw);
				xw.WriteEndElement();
			}

			xw.WriteStartElement("CallTargetInfo");
			xw.WriteAttributeString("Count", this.m_targets.Count.ToString());
			foreach(IDMapAtCallInfoSingleCallTarget idmst in this.m_targets)
				idmst.XMLExport(xw);
			xw.WriteEndElement();

			xw.WriteEndElement();
		}

		private static string BuildIDMapString(MemLocID from, HashSet<MemLocID> tolocs)
		{
			StringBuilder sb = new StringBuilder(from.ToString());

			if(tolocs.Count == 1)
			{
				sb.Append(" -> ");
				sb.Append(tolocs.First().ToString());
			}
			else
			{
				bool first = true;
				sb.Append(" -> {");
				foreach(MemLocID mm in tolocs)
				{
					if(!first)
						sb.Append(", ");
					first = false;

					sb.Append(mm.ToString());
				}
				sb.Append("}");
			}
			return sb.ToString();
		}
	}

	/// <summary>
	/// A call site with a virtual call may have many possible call targets this class captures the information for 
	/// a single target.
	/// </summary>
	public class IDMapAtCallInfoSingleCallTarget
	{
		/// <summary>
		/// The target method that we resolve as being called. 
		/// Null if the bytecode is an allocation and this is an array allocation.
		/// </summary>
		public readonly IRBodyDef CalleeTarget;

		/// <summary>
		/// The ContextID of the callee.
		/// </summary>
		public readonly int CalleeContextID;

		/// <summary>
		/// The heap state at callee method entry.
		/// </summary>
		public readonly PIModel CalleeEntryState;

		/// <summary>
		/// The heap state at callee exit.
		/// </summary>
		public readonly PIModel CalleeExitState;

		////////////////////////////

		/// <summary>
		/// Map from MemLocID values in the caller to the values in the callee.
		/// </summary>
		public readonly Dictionary<MemLocID, HashSet<MemLocID>> FromCallerToCallee;

		/// <summary>
		/// Map from MemLocID values in the callee to the values in the caller (of particular interest are any freshly allocated regions).
		/// </summary>
		public readonly Dictionary<MemLocID, HashSet<MemLocID>> FromCalleeToCaller;

		public IDMapAtCallInfoSingleCallTarget(IRBodyDef ctarget, int ctxid, PIModel centrystate, PIModel cexitstate)
		{
			this.CalleeTarget = ctarget;
			this.CalleeContextID = ctxid;

			this.CalleeEntryState = centrystate;
			this.CalleeExitState = cexitstate;

			this.FromCallerToCallee = new Dictionary<MemLocID, HashSet<MemLocID>>();
			this.FromCalleeToCaller = new Dictionary<MemLocID, HashSet<MemLocID>>();
		}

		/// <summary>
		/// Write the info for this particular call mapping to the xml stream.
		/// </summary>
		public void XMLExport(System.Xml.XmlWriter xw)
		{
			xw.WriteStartElement("CallIDMapTrgt");
			xw.WriteAttributeString("CalleeMethod", this.CalleeTarget.InvokeID.ToString());
			xw.WriteAttributeString("CalleeContextID", this.CalleeContextID.ToString());
			xw.WriteAttributeString("IsBuiltin", this.CalleeTarget.IsBuiltin.ToString());
			if(this.CalleeTarget.IsBuiltin)
				xw.WriteAttributeString("BISemantics", this.CalleeTarget.BuiltinSemanticName);

			xw.WriteStartElement("CallerCalleeMap");
			xw.WriteAttributeString("Count", this.FromCallerToCallee.Count.ToString());
			foreach(KeyValuePair<MemLocID, HashSet<MemLocID>> ccmid in this.FromCallerToCallee)
				XMLExportBinding(xw, ccmid.Key, ccmid.Value);
			xw.WriteEndElement();

			xw.WriteStartElement("CalleeCallerMap");
			xw.WriteAttributeString("Count", this.FromCalleeToCaller.Count.ToString());
			foreach(KeyValuePair<MemLocID, HashSet<MemLocID>> ccmid in this.FromCalleeToCaller)
				XMLExportBinding(xw, ccmid.Key, ccmid.Value);
			xw.WriteEndElement();

			xw.WriteEndElement();
		}

		private static void XMLExportBinding(System.Xml.XmlWriter xw, MemLocID fromids, HashSet<MemLocID> toids)
		{
			xw.WriteStartElement("MemIDMapping");
			
			fromids.XMLExportAttributesOnly(xw);

			xw.WriteStartElement("MapTrgts");
			xw.WriteAttributeString("TrgtCount", toids.Count.ToString());
			foreach(MemLocID mlid in toids)
				mlid.XMLExport(xw);
			xw.WriteEndElement();

			xw.WriteEndElement();
		}
	}
}
